#include <iostream>
#include <fstream>
#include <memory>
#include <random>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <sys/stat.h>
#include <unistd.h>
#include <cstring>

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

struct RGB {
    uint8_t r, g, b;
};

enum class ErrorCode {
    Success,
    InvalidArguments,
    FileOpenFailed,
    FileSizeMismatch,
    DirectoryCreationFailed,
    DirectoryNotWritable,
    ImageSaveFailed
};

struct NV12Data {
    std::unique_ptr<uint8_t[]> data;
    int size;
};

struct CropRegion {
    int x;
    int y;
};

// 打印用法信息
void printUsage(const char* programName) {
    std::cerr << "用法: " << programName << " <输入文件> <宽度> <高度> [裁剪宽度] [裁剪高度]\n"
              << "示例: " << programName << " test.yuv 640 640 300 300\n"
              << "如果未提供裁剪尺寸，将使用默认值300x300\n";
}

// 从文件读取 NV12 数据
ErrorCode readNV12FromFile(const std::string& file_path, int width, int height, NV12Data& out_data) {
    out_data.size = height * width * 3 / 2;
    out_data.data = std::make_unique<uint8_t[]>(out_data.size);
    
    std::ifstream file(file_path, std::ios::binary);
    if (!file) {
        return ErrorCode::FileOpenFailed;
    }
    
    file.read(reinterpret_cast<char*>(out_data.data.get()), out_data.size);
    if (file.gcount() != out_data.size) {
        return ErrorCode::FileSizeMismatch;
    }
    
    return ErrorCode::Success;
}

// 裁剪 NV12 图像
ErrorCode cropNV12(const NV12Data& nv12_data, int width, int height, 
                 int x, int y, int crop_w, int crop_h,
                 NV12Data& out_data) {
    int y_size = width * height;

    out_data.size = crop_w * crop_h * 3 / 2;
    out_data.data = std::make_unique<uint8_t[]>(out_data.size);

    // 裁剪 Y 分量
    for (int i = 0; i < crop_h; ++i) {
        int src_row = (y + i) * width + x;
        int dst_row = i * crop_w;
        memcpy(out_data.data.get() + dst_row, nv12_data.data.get() + src_row, crop_w);
    }

    // 裁剪 UV 分量
    int uv_y = y / 2;
    int uv_h = crop_h / 2;
    uint8_t* cropped_uv = out_data.data.get() + crop_w * crop_h;

    for (int i = 0; i < uv_h; ++i) {
        int src_row = y_size + (uv_y + i) * width + x;
        int dst_row = i * crop_w;
        memcpy(cropped_uv + dst_row, nv12_data.data.get() + src_row, crop_w);
    }

    return ErrorCode::Success;
}

// YUV 到 RGB 转换
void yuvToRGB(uint8_t y, uint8_t u, uint8_t v, RGB& out_rgb) {
    float y_f = y;
    float u_f = u - 128.0f;
    float v_f = v - 128.0f;
    
    float r = y_f + 1.402f * v_f;
    float g = y_f - 0.344136f * u_f - 0.714136f * v_f;
    float b = y_f + 1.772f * u_f;
    
    r = std::max(0.0f, std::min(255.0f, r));
    g = std::max(0.0f, std::min(255.0f, g));
    b = std::max(0.0f, std::min(255.0f, b));
    
    out_rgb.r = static_cast<uint8_t>(r);
    out_rgb.g = static_cast<uint8_t>(g);
    out_rgb.b = static_cast<uint8_t>(b);
}

// NV12 转 RGB
std::unique_ptr<RGB[]> nv12ToRGB(const NV12Data& nv12_data, int width, int height) {
    int y_size = width * height;
    auto rgb_image = std::make_unique<RGB[]>(y_size);
    
    const uint8_t* y_plane = nv12_data.data.get();
    const uint8_t* uv_plane = nv12_data.data.get() + y_size;
    
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            uint8_t y_val = y_plane[y * width + x];
            
            int uv_x = x / 2;
            int uv_y = y / 2;
            uint8_t u_val = uv_plane[uv_y * width + (uv_x * 2)];
            uint8_t v_val = uv_plane[uv_y * width + (uv_x * 2) + 1];
            
            yuvToRGB(y_val, u_val, v_val, rgb_image[y * width + x]);
        }
    }
    
    return rgb_image;
}

// 生成随机裁剪区域
CropRegion generateRandomCropRegion(int img_width, int img_height, int crop_w, int crop_h) {
    std::random_device rd;
    std::mt19937 gen(rd());
    
    std::uniform_int_distribution<> dist_x(0, img_width - crop_w);
    std::uniform_int_distribution<> dist_y(0, img_height - crop_h);
    
    CropRegion region;
    region.x = dist_x(gen);
    region.y = dist_y(gen);
    
    region.x = region.x & ~1;
    region.y = region.y & ~1;
    
    return region;
}

// 保存 RGB 图像为 JPG
ErrorCode saveRGBToJPG(const RGB* rgb_data, int width, int height, 
                     const std::string& output_path, int quality = 90) {
    std::unique_ptr<uint8_t[]> rgb_buffer = std::make_unique<uint8_t[]>(width * height * 3);
    for (int i = 0; i < width * height; ++i) {
        rgb_buffer[i*3]   = rgb_data[i].r;
        rgb_buffer[i*3+1] = rgb_data[i].g;
        rgb_buffer[i*3+2] = rgb_data[i].b;
    }
    
    if (!stbi_write_jpg(output_path.c_str(), width, height, 3, rgb_buffer.get(), quality)) {
        return ErrorCode::ImageSaveFailed;
    }
    
    return ErrorCode::Success;
}

int main(int argc, char* argv[]) {
    // 检查参数数量
    if (argc < 4) {
        printUsage(argv[0]);
        return static_cast<int>(ErrorCode::InvalidArguments);
    }

    // 解析命令行参数
    const std::string file_path = argv[1];
    int width = 0, height = 0;
    try {
        width = std::stoi(argv[2]);
        height = std::stoi(argv[3]);
    } catch (...) {
        std::cerr << "错误: 宽度和高度必须是整数\n";
        printUsage(argv[0]);
        return static_cast<int>(ErrorCode::InvalidArguments);
    }

    // 设置裁剪尺寸 (可选参数)
    int crop_w = 300;
    int crop_h = 300;
    if (argc >= 5) {
        try {
            crop_w = std::stoi(argv[4]);
        } catch (...) {
            std::cerr << "警告: 裁剪宽度不是有效整数，使用默认值300\n";
        }
    }
    if (argc >= 6) {
        try {
            crop_h = std::stoi(argv[5]);
        } catch (...) {
            std::cerr << "警告: 裁剪高度不是有效整数，使用默认值300\n";
        }
    }

    // 检查输入文件
    if (access(file_path.c_str(), F_OK) == -1) {
        std::cerr << "错误: 输入文件不存在: " << file_path << std::endl;
        return static_cast<int>(ErrorCode::FileOpenFailed);
    }

    // 检查输出目录
    const std::string output_dir = "test";
    struct stat st;
    if (stat(output_dir.c_str(), &st) == -1) {
        if (mkdir(output_dir.c_str(), 0755) == -1) {
            std::cerr << "错误: 无法创建输出目录: " << output_dir << std::endl;
            return static_cast<int>(ErrorCode::DirectoryCreationFailed);
        }
    }

    if (access(output_dir.c_str(), W_OK) == -1) {
        std::cerr << "错误: 输出目录不可写: " << output_dir << std::endl;
        return static_cast<int>(ErrorCode::DirectoryNotWritable);
    }

    // 读取 NV12 文件
    std::cout << "正在读取 NV12 文件..." << std::endl;
    NV12Data nv12_data;
    if (auto err = readNV12FromFile(file_path, width, height, nv12_data); err != ErrorCode::Success) {
        std::cerr << "错误: 读取 NV12 文件失败" << std::endl;
        return static_cast<int>(err);
    }
    std::cout << "NV12 文件读取成功，大小: " << nv12_data.size << " 字节" << std::endl;
    
    // 1. 保存整张图片
    std::cout << "\n正在处理整张图片..." << std::endl;
    auto full_rgb = nv12ToRGB(nv12_data, width, height);
    if (auto err = saveRGBToJPG(full_rgb.get(), width, height, output_dir + "/full_image.jpg"); 
        err != ErrorCode::Success) {
        std::cerr << "错误: 保存完整图片失败" << std::endl;
        return static_cast<int>(err);
    }
    
    // 2. 处理裁剪图片
    std::cout << "\n正在处理裁剪图片..." << std::endl;
    // 生成随机裁剪区域
    std::cout << "生成随机裁剪区域..." << std::endl;
    auto region = generateRandomCropRegion(width, height, crop_w, crop_h);
    int x = std::max(0, region.x & ~1);
    int y = std::max(0, region.y & ~1);
    crop_w = std::max(2, (crop_w + 1) & ~1);
    crop_h = std::max(2, (crop_h + 1) & ~1);
    crop_w = std::min(crop_w, width - x);
    crop_h = std::min(crop_h, height - y);
    crop_w = crop_w & ~1;
    crop_h = crop_h & ~1;

    std::cout << "调整后的裁剪参数: x=" << x << ", y=" << y 
              << ", crop_w=" << crop_w << ", crop_h=" << crop_h << std::endl;
    
    // 裁剪图像
    std::cout << "裁剪图像中..." << std::endl;
    NV12Data cropped_nv12;
    if (auto err = cropNV12(nv12_data, width, height, x, y, crop_w, crop_h, cropped_nv12); 
        err != ErrorCode::Success) {
        std::cerr << "错误: 图像裁剪失败" << std::endl;
        return static_cast<int>(err);
    }
    std::cout << "图像裁剪成功，裁剪后大小: " << cropped_nv12.size << " 字节" << std::endl;
    
    // 转换为 RGB
    std::cout << "转换为 RGB 格式..." << std::endl;
    auto cropped_rgb = nv12ToRGB(cropped_nv12, crop_w, crop_h);
    std::cout << "RGB 转换完成，像素数: " << crop_w * crop_h << std::endl;
    
    // 保存裁剪后的图片
    if (auto err = saveRGBToJPG(cropped_rgb.get(), crop_w, crop_h, output_dir + "/crop.jpg"); 
        err != ErrorCode::Success) {
        std::cerr << "错误: 保存裁剪图片失败" << std::endl;
        return static_cast<int>(err);
    }
    
    return static_cast<int>(ErrorCode::Success);
}